/**
 * Copyright (C) 2025 pony working team
 * SPDX-License-Identifier: MulanPSL-2.0
 */

/* includes (local) ----------------------------------------------------------*/
#include <hal_mtd.h>
/* includes (standard library, system) ---------------------------------------*/
#include <string.h>
#include <stdlib.h>
#include <stdbool.h>
/* includes (other library) --------------------------------------------------*/
/* includes (project) --------------------------------------------------------*/
#include <x_errno.h>

/* defines -------------------------------------------------------------------*/
/* typedefs ------------------------------------------------------------------*/
/* macro ---------------------------------------------------------------------*/
/* variables (extern) --------------------------------------------------------*/
/* variables (local) ---------------------------------------------------------*/
/* variables (global) --------------------------------------------------------*/
/* inlines -------------------------------------------------------------------*/
/* functions -----------------------------------------------------------------*/

hal_mtd_t *hal_mtd_open(const char *name)
{
    return container_of(hal_open(name, 0), hal_mtd_t, parent);
}

int hal_mtd_close(hal_mtd_t *mtd)
{
    return hal_close((hal_dev_t *)mtd);
}

int hal_mtd_erase_block(hal_mtd_t *mtd, size_t startblock, size_t nblocks)
{
    int ret;

    if (mtd == NULL || nblocks == 0) {
        return -EINVAL;
    }
    if (mtd->methods->erase_block == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    ret = mtd->methods->erase_block(mtd, startblock, nblocks);
    k_mutex_unlock(&mtd->lock);

    return ret;
}

int hal_mtd_read_block(hal_mtd_t *mtd, size_t startblock, size_t nblocks, uint8_t *buf)
{
    int ret;

    if (mtd == NULL || buf == NULL || nblocks == 0) {
        return -EINVAL;
    }
    if (mtd->methods->read_block == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    ret = mtd->methods->read_block(mtd, startblock, nblocks, buf);
    k_mutex_unlock(&mtd->lock);

    return ret;
}

int hal_mtd_write_block(hal_mtd_t *mtd, size_t startblock, size_t nblocks, const uint8_t *buf)
{
    int ret;

    if (mtd == NULL || buf == NULL || nblocks == 0) {
        return -EINVAL;
    }
    if (mtd->methods->write_block == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    ret = mtd->methods->write_block(mtd, startblock, nblocks, buf);
    k_mutex_unlock(&mtd->lock);

    return ret;
}

int hal_mtd_erase(hal_mtd_t *mtd, size_t pos, size_t nbytes)
{
    int ret;

    if (mtd == NULL || nbytes == 0) {
        return -EINVAL;
    }
    if (pos < mtd->start_addr || pos - mtd->start_addr + nbytes > mtd->total_size) {
        return -EPERM;
    }
    if (mtd->methods->erase == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    ret = mtd->methods->erase(mtd, pos, nbytes);
    k_mutex_unlock(&mtd->lock);

    return ret;
}

int hal_mtd_read(hal_mtd_t *mtd, size_t pos, size_t nbytes, uint8_t *buf)
{
    int ret;

    if (mtd == NULL || buf == NULL || nbytes == 0) {
        return -EINVAL;
    }
    if (pos < mtd->start_addr || pos - mtd->start_addr + nbytes > mtd->total_size) {
        return -EPERM;
    }
    if (mtd->methods->read == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    ret = mtd->methods->read(mtd, pos, nbytes, buf);
    k_mutex_unlock(&mtd->lock);

    return ret;
}

static int __mtd_write(hal_mtd_t *mtd, size_t pos, size_t nbytes, const uint8_t *buf)
{
    size_t page_size = mtd->page_size;
    size_t cur_size = page_size - (pos % page_size);
    size_t cur_addr = pos;
    size_t end_addr = pos + nbytes;

    do {
        if (mtd->methods->write(mtd, cur_addr, cur_size, buf) != 0) {
            return -1;
        }
        /* Update the address and size variables for next page programming */
        cur_addr += cur_size;
        buf += cur_size;
        cur_size = ((cur_addr + page_size) > end_addr) ? end_addr - cur_addr : page_size;
    } while (cur_addr < end_addr);

    return 0;
}

int hal_mtd_write(hal_mtd_t *mtd, size_t pos, size_t nbytes, const uint8_t *buf)
{
    int ret;

    if (mtd == NULL || buf == NULL || nbytes == 0) {
        return -EINVAL;
    }
    if (pos < mtd->start_addr || pos - mtd->start_addr + nbytes > mtd->total_size) {
        return -EPERM;
    }
    if (mtd->methods->write == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    if (mtd->page_size <= 1) {
        ret = mtd->methods->write(mtd, pos, nbytes, buf);
    } else {
        /* Perform the write page by page */
        ret = __mtd_write(mtd, pos, nbytes, buf);
    }

    k_mutex_unlock(&mtd->lock);

    return ret;
}

static int __mtd_ioctl(hal_mtd_t *mtd, int cmd, ...)
{
    va_list valist;
    int ret;

    va_start(valist, cmd);
    ret = mtd->methods->ioctl(mtd, cmd, valist);
    va_end(valist);

    return ret;
}

static int __mtd_erase_write(hal_mtd_t *mtd, size_t pos, size_t nbytes, const uint8_t *buf)
{
    if (mtd->methods->erase(mtd, pos, nbytes) != 0) {
        return -1;
    }
    return mtd->methods->write(mtd, pos, nbytes, buf);
}

int hal_mtd_write_safe(hal_mtd_t *mtd, size_t pos, size_t nbytes, const uint8_t *buf)
{
    int ret = 0;
    uint8_t *block_buf = NULL;
    size_t sector_size = 0;
    size_t sector_offset = 0;
    size_t write_len = 0; /* length of each write */
    const hal_mtd_methods_t *methods = mtd->methods;

    if (mtd == NULL || buf == NULL || nbytes == 0) {
        return -EINVAL;
    }
    if (pos < mtd->start_addr || pos - mtd->start_addr + nbytes > mtd->total_size) {
        return -EPERM;
    }
    if (methods->erase == NULL || methods->write == NULL || methods->read == NULL ||
        methods->ioctl == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }

    do {
        /*
         * note: the flash sector size must be reacquired each time
         *  1) the sector size of inner flash is unfixed generally.
         *  2) the sector size of spi flash is fixed
         */
        if (__mtd_ioctl(mtd, HAL_MTD_CMD_GET_SECTOR_SIZE, pos, &sector_size) != 0) {
            return -1;
        }

        /**
         *       .-- aligned_pos                              .--- 1)          .-- 2)
         *      |                  pos -.                    /                /
         *      |                        \  .----------- nbytes --------------.
         *      |                         |/                   `.              \
         *   ~~~+=========================@=====================@===+~~...~~+===@~~~~
         *   ///|/////////////////////////|                         |       |   |\\\\
         *   ~~~+=========================@=========================+~~...~~+===@~~~~
         *      |                         |                         |       |
         *      |`----- sector_offset ---` `----- remain_len ------`|`--#--`
         *       \                                                 /    \
         *        `-------------------- sector_size --------------`      `--- (n * sector_size)
         */

        sector_offset = pos & (sector_size - 1);
        write_len = min(nbytes, sector_size - sector_offset);

        if (sector_offset == 0 && nbytes >= sector_size) {   // override one sector
            if ((ret = __mtd_erase_write(mtd, pos, write_len, buf)) < 0) {
                goto EXIT_FREE;
            }
        } else {
            size_t aligned_pos = pos - sector_offset;
            block_buf = (uint8_t *)malloc(sector_size);
            if (block_buf == NULL) {
                ret = -ENOMEM;
                goto EXIT_UNLOCK;
            }
            if ((ret = methods->read(mtd, aligned_pos, sector_size, block_buf)) != 0) {
                goto EXIT_FREE;
            }
            memcpy(block_buf + sector_offset, buf, write_len);
            if ((ret = __mtd_erase_write(mtd, aligned_pos, sector_size, block_buf)) != 0) {
                goto EXIT_FREE;
            }
            if (block_buf != NULL) {
                free(block_buf);
                block_buf = NULL;
            }
        }
        nbytes -= write_len;
        buf += write_len;
        pos += write_len;
    } while (nbytes > 0);

    ret = 0;

EXIT_FREE:
    if (block_buf != NULL) {
        free(block_buf);
    }
EXIT_UNLOCK:
    k_mutex_unlock(&mtd->lock);

    return ret;
}

int hal_mtd_erase_write(hal_mtd_t *mtd, size_t pos, size_t nbytes, const uint8_t *buf)
{
    int ret;

    if (mtd == NULL || buf == NULL || nbytes == 0) {
        return -EINVAL;
    }
    if (pos < mtd->start_addr || pos - mtd->start_addr + nbytes > mtd->total_size) {
        return -EPERM;
    }
    if (mtd->methods->erase == NULL || mtd->methods->write == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    ret = __mtd_erase_write(mtd, pos, nbytes, buf);
    k_mutex_unlock(&mtd->lock);

    return ret;
}

int hal_mtd_ioctl(hal_mtd_t *mtd, int cmd, ...)
{
    va_list valist;
    int ret;

    if (mtd == NULL) {
        return -EINVAL;
    }
    if (mtd->methods->ioctl == NULL) {
        return -ENOSYS;
    }
    if (unlikely(k_mutex_lock(&mtd->lock) != 0)) {
        return -EBUSY;
    }
    va_start(valist, cmd);
    ret = mtd->methods->ioctl(mtd, cmd, valist);
    va_end(valist);
    k_mutex_unlock(&mtd->lock);

    return ret;
}

int hal_mtd_register(hal_mtd_t *mtd, const char *name, const hal_mtd_methods_t *methods,
                     uintptr_t data)
{
    if (mtd == NULL || methods == NULL) {
        return -EINVAL;
    }
    mtd->parent.name = name;
    mtd->methods = methods;

    if (mtd->page_size == 0) {
        mtd->page_size = 1;
    }

    if (k_mutex_init(&mtd->lock) != 0) {
        return -ENOMEM;
    }

    if (hal_register(&mtd->parent) != 0) {
        k_mutex_destroy(&mtd->lock);
        return -1;
    }

    return 0;
}
